import fitz
from typing import Optional, Union, Callable, List, Tuple, Any

from app.utils.pymupdf_tools import pdf_get_content_region
from tests.base_test import base_test_case

logger = base_test_case.get_logger(__name__)
TEST_DATA_DIR = base_test_case.test_data_dir
OUTPUT_DATA_DIR = base_test_case.output_data_dir
RectLike = Union[None, Tuple[float, float, float, float], fitz.Rect]

# 定义颜色和边框宽度常量
COLOR_RED = fitz.pdfcolor["red"]
COLOR_GREEN = fitz.pdfcolor["green"]
BORDER_WIDTH = 0.3
# pdf表格数据结构类型映射
DATA_STRUCTURE_TYPES_MAP = {
    'html': lambda df: df.to_html(escape=False, border=0, classes=None, bold_rows=False),
    'dict': lambda df: df.to_dict(orient='dict'),
    'list': lambda df: df.to_dict(orient='list'),
    'records': lambda df: df.to_dict(orient='records'),
    'index': lambda df: df.to_dict(orient='index'),
}


def extract_tables_from_pdf(
        pdf_path: str,
        page_range: Optional[Union[int, List[int], Tuple[int, int]]] = None,
        clip: Union[RectLike, Callable[[fitz.Page], RectLike]] = None,
        strategy: str = "lines_strict",
        vertical_strategy: str = "lines_strict",
        horizontal_strategy: str = "lines_strict",
        min_words_vertical: int = 1,
        min_words_horizontal: int = 1,
        intersection_tolerance: int = 3,
        data_structure_type: str = "html",
):
    """
    从 PDF 中提取表格内容，支持灵活的页面选择、区域裁剪（包括函数动态计算）和表格识别策略。

    参数:
        pdf_path (str): PDF 文件路径。
        page_range (int, list, tuple, None): 要处理的页码。
            - int: 单页，如 0
            - list: 多页，如 [0, 2, 4]
            - tuple: 页码范围（包含两端），如 (0, 3) 表示第 0,1,2,3 页
            - None: 所有页
        clip: 裁剪区域，可以是：
            - None：整页
            - (x0, y0, x1, y1)：固定矩形
            - fitz.Rect：
            - callable(page) -> RectLike：每页动态计算 ROI
        strategy (str): 主策略，可选 "lines", "lines_strict", "text", "explicit"
            - lines 有边框表格. 基于页面中显式绘制的水平/垂直线构建网格。容忍一定交叉误差（由 intersection_tolerance 控制）。 适用于 表格边框断裂、虚线、或颜色很浅	先尝试 "lines" + 增大 intersection_tolerance（如 5~10)
            - lines_strict 标准带边框表格（线条清晰、对齐良好）.比 "lines" 更严格：要求线条必须形成闭合单元格，交叉点必须精确匹配。适合印刷体 PDF 表格。 适用于 表格有完整边框（黑线/灰线）
            - text 无边框表格（纯靠文本对齐排版）. 不依赖线条，而是通过分析文本块的 x/y 坐标聚类，推断行列结构。适用于讲义、教材、财报等“视觉表格”。 适用于 表格无边框，但列对齐整齐（如 Word 导出
            - explicit 极简模式 / 调试用,仅当表格区域被显式标记为“表格”（如 Tagged PDF 中的 <Table> 标签）时才识别。绝大多数普通 PDF 不适用！
        vertical_strategy (str): 垂直线检测策略
        horizontal_strategy (str): 水平线检测策略
        min_words_vertical (int): 垂直方向构成一列所需的最少单词数
        min_words_horizontal (int): 水平方向构成一行所需的最少单词数
        intersection_tolerance (int): 线条交叉点容差（像素）
        data_structure_type: 表格数据的返回类型，类型：dict，list，html；默认为 html

    返回:
        list[dict]: 每个元素包含 'page', 'table_index', 'data'（二维列表）
    """
    doc = fitz.open(pdf_path)
    all_tables = []
    total_pages = len(doc)

    # 解析 page_range
    if page_range is None:
        pages_to_process = list(range(total_pages))
    elif isinstance(page_range, int):
        pages_to_process = [page_range]
    elif isinstance(page_range, list):
        pages_to_process = page_range
    elif isinstance(page_range, tuple) and len(page_range) == 2:
        start, end = page_range
        if start > end:
            raise ValueError("In tuple page_range, start must be <= end.")
        pages_to_process = list(range(start, end + 1))
    else:
        raise ValueError("Invalid page_range format.")

    print(f"[INFO] Processing pages (1-based): {[p + 1 for p in pages_to_process]}")
    # 判断 clip 是否为函数
    is_clip_callable = callable(clip)

    for page_num in pages_to_process:
        if not (0 <= page_num < total_pages):
            print(f"Warning: Page {page_num} is out of range (total pages: {total_pages}). Skipped.")
            continue

        page = doc.load_page(page_num)

        # 动态计算 clip 区域（如果是函数）
        if is_clip_callable:
            try:
                current_clip = clip(page)  # 用户函数返回 Rect 或 tuple
            except Exception as e:
                print(f"Error calling clip function on page {page_num}: {e}")
                current_clip = None
        else:
            current_clip = clip

        # 标准化为 fitz.Rect 或 None
        clip_rect = fitz.Rect(current_clip) if current_clip is not None else None

        # 表格识别
        tabs = page.find_tables(
            clip=clip_rect,
            strategy=strategy,
            vertical_strategy=vertical_strategy,
            horizontal_strategy=horizontal_strategy,
            min_words_vertical=min_words_vertical,
            min_words_horizontal=min_words_horizontal,
            intersection_tolerance=intersection_tolerance,
        )

        if tabs.tables:
            print(f"Page {page_num + 1} has {len(tabs.tables)} table(s).")
            for tab_idx, tab in enumerate(tabs.tables):
                tab_df = tab.to_pandas()
                bbox_cells = tab.cells
                bbox_tab = [int(bbox_cells[0][0] - 1),
                            int(bbox_cells[0][1] - 1),
                            # 右下角
                            int(max([bbox[2] for bbox in bbox_cells])),
                            # 右下角
                            int(max([bbox[-1] for bbox in bbox_cells]))  # 获取每个元组的最后一个元素并形成列表,找出最大值
                            ]

                table_data = tab.extract()
                all_tables.append({
                    "page": page_num + 1,
                    "table_index": tab_idx + 1,
                    "data": table_data,
                    'bbox': bbox_tab,
                    'count_row': tab.row_count,
                    'count_col': tab.col_count,
                    'text': DATA_STRUCTURE_TYPES_MAP[data_structure_type](tab_df),
                })
        else:
            print(f"No tables found on page {page_num + 1}")

    doc.close()
    return all_tables


if __name__ == '__main__':
    """表格提取方法："""
    input_file = str(TEST_DATA_DIR / "table.jpg")
    pdf_path = str(TEST_DATA_DIR / "25-注会-轻1-财务成本管理[上册](第3章).pdf")


    # # 获取正文区域
    # clip = pdf_get_content_region(page, zoom_factor=1.0)
    # text_dict = page.get_text('dict', clip=clip, sort=True)

    # 自定义 ROI 检测器
    def my_roi_detector(page: fitz.Page) -> fitz.Rect:
        # 示例：简单去掉上下 50pt 边距
        rect = pdf_get_content_region(page, zoom_factor=1.0)
        return rect

        # 或者调用你的实际函数：
        # return pdf_get_content_region(page, zoom_factor=1.0, is_show_log=False)


    # 自定义表格识别
    # data_structure_type: 表格数据的返回类型，类型：dict，list，html
    tables = extract_tables_from_pdf(
        pdf_path=pdf_path,
        clip=my_roi_detector,
        page_range=(0, 0),
        strategy="lines_strict",
        vertical_strategy="lines_strict",
        horizontal_strategy="lines_strict",
        data_structure_type="html",
        # min_words_vertical=3,
        # min_words_horizontal=3,
        # intersection_tolerance=1,
    )

    print(tables)
    print("Extracted tables:")
    for t in tables:
        print(f"  Page {t['page']}, Table {t['table_index']}: {len(t['data'])} rows")
    pass
